`timescale 1ns / 1ps

module sm4_top_tb;
    reg clk;
    reg reset_n;
    reg sm4_enable_in;
    reg encdec_enable_in;
    reg encdec_sel_in;
    reg valid_in;
    reg [127:0] data_in;
    reg enable_key_exp_in;
    reg user_key_valid_in;
    reg [127:0] user_key_in;
    wire ready_out;
    wire [127:0] result_out;
    wire key_exp_ready_out;
    parameter CLK_PERIOD = 10;
    parameter TEST_CYCLES = 1_000_000;
    parameter [127:0] EXPECTED_CIPHER = 128'h595298C7C6FD271F0402F804C33D3F66;
    sm4_top uut (
        .clk(clk),
        .reset_n(reset_n),
        .sm4_enable_in(sm4_enable_in),
        .encdec_enable_in(encdec_enable_in),
        .encdec_sel_in(encdec_sel_in),
        .valid_in(valid_in),
        .data_in(data_in),
        .enable_key_exp_in(enable_key_exp_in),
        .user_key_valid_in(user_key_valid_in),
        .user_key_in(user_key_in),
        .key_exp_ready_out(key_exp_ready_out),
        .ready_out(ready_out),
        .result_out(result_out)
    );
    always begin
        #(CLK_PERIOD/2) clk = ~clk;
    end
    integer i;
    integer error_count = 0;
    reg [127:0] last_result;
    
    initial begin
        clk = 0;
        reset_n = 0;
        sm4_enable_in = 0;
        encdec_enable_in = 0;
        encdec_sel_in = 0;
        valid_in = 0;
        data_in = 0;
        enable_key_exp_in = 0;
        user_key_valid_in = 0;
        user_key_in = 0;
        #(CLK_PERIOD*2);
        reset_n = 1;
        @(posedge clk);
        user_key_in = 128'h0123456789ABCDEFFEDCBA9876543210;
        user_key_valid_in = 1;
        enable_key_exp_in = 1;
        wait(key_exp_ready_out);
        @(posedge clk);
        user_key_valid_in = 0;
        enable_key_exp_in = 0;
        data_in = 128'h0123456789ABCDEFFEDCBA9876543210;
        encdec_sel_in = 0;
        encdec_enable_in = 1;
        sm4_enable_in = 1;
        
        $display("Starting SM4 encryption test for %0d iterations...", TEST_CYCLES);
        $timeformat(-9, 3, " ns", 10);
        $display("Start time: %t", $realtime);
        for (i = 0; i < TEST_CYCLES; i = i + 1) begin
            @(posedge clk);
            valid_in = 1;
            @(posedge ready_out);
            valid_in = 0;
            last_result = result_out;
            if (i % 100_000 == 0) begin
                $display("Completed %0d iterations at %t", i, $realtime);
            end
            data_in = result_out;
        end
        if (last_result === EXPECTED_CIPHER) begin
            $display("Test PASSED at %t. Final ciphertext matches expected value: %h", $realtime, last_result);
        end else begin
            $display("Test FAILED at %t. Expected: %h, Got: %h", $realtime, EXPECTED_CIPHER, last_result);
            error_count = error_count + 1;
        end
        $display("Test completed at %t", $realtime);
        $display("Total iterations: %0d", TEST_CYCLES);
        $display("Error count: %0d", error_count);
        $display("Last ciphertext: %h", last_result);
        
        $finish;
    end
    
endmodule